"""
Domain Limits
=============
"""

import numpy

from chaco.plot import Plot, ArrayPlotData
from chaco.api import ToolbarPlot
from chaco.tools.api import PanTool, ZoomTool
from enable.api import ComponentEditor
from traits.api import Instance, HasTraits
from traitsui.api import View, Item

"""
Plot where depth is the index such that the plot is vertical and the origin is 
the upper left
"""
import numpy
from chaco.api import ToolbarPlot, ArrayPlotData
from chaco.tools.api import LineInspector
from enable.api import ComponentEditor
from traits.api import HasTraits, Instance
from traitsui.api import UItem, View

"""
Controlling aspect ratio

Demonstrates various ways that Chaco can control aspect ratios in plots.

Left-drag pans the plot.

Mousewheel up and down zooms the plot in and out.
"""

# Major library imports
from itertools import chain
from numpy import sort, isfinite
from numpy.random import random

# Enthought library imports
from enable.api import ComponentEditor
from traits.api import *
from traitsui.api import *

# Chaco imports
from chaco.api import ArrayPlotData, Plot, create_line_plot,OverlayPlotContainer,HPlotContainer,Plot,ArrayPlotData,viridis
from chaco.tools.api import PanTool, ZoomTool
from chaco.tools.cursor_tool import CursorTool, BaseCursorTool


if __name__ == "__main__":
    import data_cube
    import edit_line
    import functionplotter
    import hyetograph
    import multi_line_plot_demo
    import nonlinear_color_mapping
    import data_stream
else:
    from . import data_cube
    from . import edit_line
    from . import functionplotter
    from . import hyetograph
    from . import multi_line_plot_demo
    from . import nonlinear_color_mapping
    from . import data_stream


class AspectRatio(BaseFloat):
    "A new Trait for defining aspect ratios"

    default_value = 1.0

    info_text = "a nonzero floating point number"

    def validate(self, object, name, value):
        value = super(AspectRatio, self).validate(object, name, value)
        if value != 0.0 and isfinite(value):
            return value
        else:
            self.error(object, name, value)


class MyPlot(HasTraits):
    plot = Any()
    screen_enabled = Bool(False)
    screen_aspect = AspectRatio()
    fixed_x = Bool(False)
    fixed_y = Bool(False)
    traits_view = View(
        VGroup(
            HGroup(
                Item("screen_enabled", label="Screen"),
                Item("screen_aspect", label="aspect ratio (w/h)"),
            ),
            HGroup(
                Item("fixed_x", label="Data X fixed"),
                Item("fixed_y", label="Data Y fixed"),
            ),
            Item(
                "plot",
                editor=ComponentEditor(size=(100, 100)),
                show_label=False,
            ),
        ),
        width=600,
        height=600,
        resizable=True,
        title="Aspect Ratio Example",
    )

    def _plot_default(self):
        numpoints = 200
        plotdata = ArrayPlotData(
            x=sort(random(numpoints)), y=random(numpoints)
        )
        plot = Plot(plotdata)
        plot.plot(("x", "y"), type="scatter")
        plot.tools.append(PanTool(plot))
        plot.overlays.append(ZoomTool(plot))
        return plot

    def _screen_enabled_changed(self):
        if self.screen_enabled:
            self.plot.aspect_ratio = self.screen_aspect
        else:
            self.plot.aspect_ratio = None
        self.plot.request_redraw()

    def _screen_aspect_changed(self):
        if self.screen_enabled:
            self.plot.aspect_ratio = self.screen_aspect
            self.plot.request_redraw()

    def _fixed_x_changed(self):
        self.plot.x_mapper.stretch_data = not self.fixed_x
        # Also have to change all the renderers' mappers
        for renderer in chain(*list(self.plot.plots.values())):
            renderer.index_mapper.stretch_data = not self.fixed_x
        self.plot.request_redraw()

    def _fixed_y_changed(self):
        self.plot.y_mapper.stretch_data = not self.fixed_y
        for renderer in chain(*list(self.plot.plots.values())):
            renderer.value_mapper.stretch_data = not self.fixed_y
        self.plot.request_redraw()


class MyPlot1(HasTraits):
    """Plot where depth is the index such that the plot is vertical
    and the origin is the upper left
    """

    plot = Instance(ToolbarPlot)

    traits_view = View(
        UItem("plot", editor=ComponentEditor()),
        width=600,
        height=600,
        resizable=True,
    )

    def _plot_default(self):
        depth = numpy.arange(1.0, 100.0, 0.1)
        data_series = numpy.sin(depth) + depth / 10.0
        plot_data = ArrayPlotData(index=depth)
        plot_data.set_data("data_series", data_series)
        plot = ToolbarPlot(plot_data, orientation="v", origin="top left")
        line = plot.plot(("index", "data_series"))[0]

        line_inspector = LineInspector(component=line, write_metadata=True)
        line.tools.append(line_inspector)
        line.overlays.append(line_inspector)

        return plot


class ExamplePlotApp(HasTraits):
    plot = Instance(Plot)

    def _plot_default(self):
        index = numpy.arange(1.0, 10.0, 0.01)
        series1 = (100.0 + index) / (100.0 - 20 * index ** 2 + 5.0 * index ** 4)
        series2 = (100.0 + index) / (100.0 - 20 * index ** 2 + 5.0 * index ** 3)

        plot_data = ArrayPlotData(index=index)
        plot_data.set_data("series1", series1)
        plot_data.set_data("series2", series2)

        plot = ToolbarPlot(plot_data)
        line_plot = plot.plot(("index", "series1"), color="auto")[0]

        # Add pan and zoom tools
        line_plot.tools.append(PanTool(line_plot))
        line_plot.tools.append(ZoomTool(line_plot))

        # Set the domain_limits
        line_plot.index_mapper.domain_limits = (3.3, 6.6)

        return plot

    traits_view = View(
        Item(
            "plot",
            editor=ComponentEditor(),
            show_label=False,
        ),
        width=600,
        height=600,
        resizable=True,
    )

class CursorTest(HasTraits):
    plot = Instance(HPlotContainer)
    cursor1 = Instance(BaseCursorTool)
    cursor2 = Instance(BaseCursorTool)

    cursor1pos = DelegatesTo("cursor1", prefix="current_position")
    cursor2pos = DelegatesTo("cursor2", prefix="current_position")

    def __init__(self):
        # The delegates views don't work unless we caller the superclass __init__
        super(CursorTest, self).__init__()

        container = HPlotContainer(padding=0, spacing=20)
        self.plot = container
        # a subcontainer for the first plot.
        # I'm not sure why this is required. Without it, the layout doesn't work right.
        subcontainer = OverlayPlotContainer(padding=40)
        container.add(subcontainer)

        # make some data
        index = numpy.linspace(-10, 10, 512)
        value = numpy.sin(index)

        # create a LinePlot instance and add it to the subcontainer
        line = create_line_plot(
            [index, value],
            add_grid=True,
            add_axis=True,
            index_sort="ascending",
            orientation="h",
        )
        subcontainer.add(line)

        # here's our first cursor.
        csr = CursorTool(line, drag_button="left", color="blue")
        self.cursor1 = csr
        # and set it's initial position (in data-space units)
        csr.current_position = 0.0, 0.0

        # this is a rendered component so it goes in the overlays list
        line.overlays.append(csr)

        # some other standard tools
        line.tools.append(PanTool(line, drag_button="right"))
        line.overlays.append(ZoomTool(line))

        # make some 2D data for a colourmap plot
        xy_range = (-5, 5)
        x = numpy.linspace(xy_range[0], xy_range[1], 100)
        y = numpy.linspace(xy_range[0], xy_range[1], 100)
        X, Y = numpy.meshgrid(x, y)
        Z = numpy.sin(X) * numpy.arctan2(Y, X)

        # easiest way to get a CMapImagePlot is to use the Plot class
        ds = ArrayPlotData()
        ds.set_data("img", Z)

        img = Plot(ds, padding=40)
        cmapImgPlot = img.img_plot(
            "img", xbounds=xy_range, ybounds=xy_range, colormap=viridis
        )[0]

        container.add(img)

        # now make another cursor
        csr2 = CursorTool(
            cmapImgPlot, drag_button="left", color="white", line_width=2.0
        )
        self.cursor2 = csr2

        csr2.current_position = 1.0, 1.5

        cmapImgPlot.overlays.append(csr2)

        # add some standard tools. Note, I'm assigning the PanTool to the
        # right mouse-button to avoid conflicting with the cursors
        cmapImgPlot.tools.append(PanTool(cmapImgPlot, drag_button="right"))
        cmapImgPlot.overlays.append(ZoomTool(cmapImgPlot))

    traits_view = View(
        VGroup(
            HGroup(
                Item(
                    "plot",
                    editor=ComponentEditor(),
                    resizable=True,
                    springy=True,
                    show_label=False,
                ),
                springy=True,
            ),
            HGroup(
                Item("cursor1pos", width=300), Item("cursor2pos", width=300)
            ),
        ),
        title="Cursor Tool Demo",
        resizable=True,
        width=800,
        height=420,
    )





class Form(HasTraits):
    myplot = Instance(MyPlot, ())
    myplot1 = Instance(MyPlot1, ())
    plotapp = Instance(ExamplePlotApp, ())
    CursorTest_app=Instance(CursorTest, ())
    Data_cube_app=Instance(data_cube.Demo,())
    Edit_line_app=Instance(edit_line.demo,())
    Functionplotter_app=Instance(functionplotter.demo,())
    Hyetograph_app=Instance(hyetograph.popup,())
    Multi_line_plot_app=Instance(multi_line_plot_demo.demo,())
    Nonlinear_color_app=Instance(nonlinear_color_mapping.demo,())
    Data_stream_app=Instance(data_stream.popup,())
    traits_view=View(
        Tabbed(
            UItem('myplot',style='custom'),
            UItem('myplot1',style='custom'),
            UItem('plotapp',style='custom'),
            UItem('CursorTest_app', style='custom'),
            UItem('Data_cube_app', style='custom'),
            UItem('Edit_line_app', style='custom'),
            UItem('Functionplotter_app',style='custom'),
            UItem('Hyetograph_app',style='custom'),
            UItem('Multi_line_plot_app',style='custom'),
            UItem('Nonlinear_color_app',style='custom'),
            UItem('Data_stream_app',style='custom'),
        ),
        resizable=True,
    )

if __name__ == "__main__":
    Form().configure_traits()
